/*
 * #%L
 * Alfresco Repository
 * %%
 * Copyright (C) 2005 - 2023 Alfresco Software Limited
 * %%
 * This file is part of the Alfresco software. 
 * If the software was purchased under a paid Alfresco license, the terms of 
 * the paid license agreement will prevail.  Otherwise, the software is 
 * provided under the following open source license terms:
 * 
 * Alfresco is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * Alfresco is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with Alfresco. If not, see <http://www.gnu.org/licenses/>.
 * #L%
 */
package org.alfresco.encryption;

import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.DESedeKeySpec;
import jakarta.transaction.NotSupportedException;
import jakarta.transaction.SystemException;
import jakarta.transaction.UserTransaction;

import org.apache.commons.codec.binary.Base64;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.springframework.context.ApplicationContext;

import org.alfresco.repo.security.authentication.AuthenticationUtil;
import org.alfresco.service.transaction.TransactionService;
import org.alfresco.util.ApplicationContextHelper;
import org.alfresco.util.GUID;

/**
 * 
 * @since 4.0
 *
 */
public class KeyStoreTests
{
    private static ApplicationContext ctx = ApplicationContextHelper.getApplicationContext();

    private TransactionService transactionService;
    private KeyStoreChecker keyStoreChecker;
    private EncryptionKeysRegistry encryptionKeysRegistry;
    private UserTransaction txn = null;
    private KeyResourceLoader keyResourceLoader;
    private List<String> toDelete;
    private DefaultEncryptor backupEncryptor;

    @Before
    public void setup() throws SystemException, NotSupportedException
    {
        transactionService = (TransactionService) ctx.getBean("transactionService");
        keyStoreChecker = (KeyStoreChecker) ctx.getBean("keyStoreChecker");
        encryptionKeysRegistry = (EncryptionKeysRegistry) ctx.getBean("encryptionKeysRegistry");
        keyResourceLoader = (KeyResourceLoader) ctx.getBean("springKeyResourceLoader");
        backupEncryptor = (DefaultEncryptor) ctx.getBean("backupEncryptor");

        toDelete = new ArrayList<String>(10);

        AuthenticationUtil.setRunAsUserSystem();
        UserTransaction txn = transactionService.getUserTransaction();
        txn.begin();
    }

    @After
    public void teardown() throws IllegalStateException, SecurityException, SystemException
    {
        if (txn != null)
        {
            txn.rollback();
        }

        for (String guid : toDelete)
        {
            File file = new File(guid);
            if (file.exists())
            {
                file.delete();
            }
        }
    }

    public String generateEncodedKey()
    {
        try
        {
            return Base64.encodeBase64String(generateKeyData());
        }
        catch (Throwable e)
        {
            fail("Unexpected exception: " + e.getMessage());
            return null;
        }
    }

    public byte[] generateKeyData() throws NoSuchAlgorithmException
    {
        SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
        random.setSeed(System.currentTimeMillis());
        byte bytes[] = new byte[DESedeKeySpec.DES_EDE_KEY_LEN];
        random.nextBytes(bytes);
        return bytes;
    }

    protected String generateKeystoreName()
    {
        String guid = GUID.generate();
        toDelete.add(guid);
        return guid;
    }

    protected Key generateSecretKey(String keyAlgorithm)
    {
        try
        {
            DESedeKeySpec keySpec = new DESedeKeySpec(generateKeyData());
            SecretKeyFactory kf = SecretKeyFactory.getInstance(keyAlgorithm);
            SecretKey secretKey = kf.generateSecret(keySpec);
            return secretKey;
        }
        catch (Throwable e)
        {
            fail("Unexpected exception: " + e.getMessage());
            return null;
        }
    }

    protected TestAlfrescoKeyStore getKeyStore(String name, String type, final Map<String, String> passwords, final Map<String, String> encodedKeyData,
            String keyStoreLocation, String backupKeyStoreLocation)
    {
        KeyResourceLoader testKeyResourceLoader = new KeyResourceLoader() {
            @Override
            public InputStream getKeyStore(String keyStoreLocation)
                    throws FileNotFoundException
            {
                return keyResourceLoader.getKeyStore(keyStoreLocation);
            }

            @Override
            public Properties loadKeyMetaData(String keyMetaDataFileLocation)
                    throws IOException, FileNotFoundException
            {
                Properties p = new Properties();
                p.put("keystore.password", "password");
                StringBuilder aliases = new StringBuilder();
                for (String keyAlias : passwords.keySet())
                {
                    p.put(keyAlias + ".password", passwords.get(keyAlias));
                    if (encodedKeyData != null && encodedKeyData.get(keyAlias) != null)
                    {
                        p.put(keyAlias + ".keyData", encodedKeyData.get(keyAlias));
                    }
                    p.put(keyAlias + ".algorithm", "DESede");
                    aliases.append(keyAlias);
                    aliases.append(",");
                }
                if (aliases.length() > 0)
                {
                    // remove trailing comma
                    aliases.delete(aliases.length() - 1, aliases.length());
                }
                p.put("aliases", aliases.toString());
                return p;
            }
        };

        KeyStoreParameters keyStoreParameters = new KeyStoreParameters(name, type, null, "", keyStoreLocation);
        KeyStoreParameters backupKeyStoreParameters = new KeyStoreParameters(name + ".backup", type, null, "", backupKeyStoreLocation);
        TestAlfrescoKeyStore keyStore = new TestAlfrescoKeyStore();
        keyStore.setKeyStoreParameters(keyStoreParameters);
        keyStore.setBackupKeyStoreParameters(backupKeyStoreParameters);
        keyStore.setKeyResourceLoader(testKeyResourceLoader);
        keyStore.setValidateKeyChanges(true);
        keyStore.setEncryptionKeysRegistry(encryptionKeysRegistry);
        return keyStore;
    }

    @Test
    public void test1()
    {
        // missing keystore, missing backup keystore, no registered keys -> create key store with metadata key and register key

        TestAlfrescoKeyStore missingMainKeyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
                Collections.singletonMap(KeyProvider.ALIAS_METADATA, generateEncodedKey()), generateKeystoreName(), generateKeystoreName());
        missingMainKeyStore.setKeysToValidate(new HashSet<>(Collections.singletonList("metadata")));
        encryptionKeysRegistry.unregisterKey(KeyProvider.ALIAS_METADATA);
        keyStoreChecker.setMainKeyStore(missingMainKeyStore);

        try
        {
            keyStoreChecker.validateKeyStores();
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception: " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }

        assertTrue("", encryptionKeysRegistry.getRegisteredKeys(missingMainKeyStore.getKeyAliases()).contains(KeyProvider.ALIAS_METADATA));
        assertTrue("", missingMainKeyStore.exists());
        assertTrue("", missingMainKeyStore.getKey(KeyProvider.ALIAS_METADATA) != null);
    }

    @Test
    public void test2()
    {
        // missing main keystore, missing backup keystore, metadata registered key -> error, re-instate the keystore
        TestAlfrescoKeyStore missingMainKeyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
                null, generateKeystoreName(), generateKeystoreName());
        missingMainKeyStore.setKeysToValidate(new HashSet<>(Collections.singletonList("metadata")));
        assertTrue("", encryptionKeysRegistry.isKeyRegistered("metadata"));

        keyStoreChecker.setMainKeyStore(missingMainKeyStore);

        try
        {
            keyStoreChecker.validateKeyStores();
            fail("Should have caught missing main keystore");
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            // ok, expected
        }
    }

    @Test
    public void test3()
    {
        // main keystore exists, no registered metadata key -> register key

        // create main keystore
        TestAlfrescoKeyStore mainKeyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
                null, generateKeystoreName(), generateKeystoreName());
        mainKeyStore.setKeysToValidate(new HashSet<>(Collections.singletonList("metadata")));
        createAndPopulateKeyStore(mainKeyStore);

        // de-register metadata key
        encryptionKeysRegistry.unregisterKey(KeyProvider.ALIAS_METADATA);

        // check keys
        keyStoreChecker.setMainKeyStore(mainKeyStore);

        try
        {
            keyStoreChecker.validateKeyStores();
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception: " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }

        assertTrue("", encryptionKeysRegistry.isKeyRegistered(KeyProvider.ALIAS_METADATA));
    }

    @Test
    public void test4()
    {
        // create keystore, change key -> check for exception InvalidKey

        // Firstly, create main keystore to register a well-known key

        encryptionKeysRegistry.unregisterKey(KeyProvider.ALIAS_METADATA);

        TestAlfrescoKeyStore keyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
                null, generateKeystoreName(), generateKeystoreName());
        keyStore.setKeysToValidate(new HashSet<>(Collections.singletonList("metadata")));
        createAndPopulateKeyStore(keyStore);

        keyStoreChecker.setMainKeyStore(keyStore);

        // check keys
        try
        {
            // should register the metadata key
            keyStoreChecker.validateKeyStores();
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception: " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }

        // check that the metadata key has been registered
        assertTrue("", encryptionKeysRegistry.isKeyRegistered("metadata"));

        // a changed main keystore with a different metadata key
        // TestAlfrescoKeyStore changedMainKeyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
        // null, generateKeystoreName(), generateKeystoreName());
        // createAndPopulateKeyStore(changedMainKeyStore);
        keyStore.changeKey(KeyProvider.ALIAS_METADATA, generateSecretKey("DESede"));

        // keyStoreChecker.setMainKeyStore(changedMainKeyStore);
        try
        {
            keyStoreChecker.validateKeyStores();
            fail("Expected key store checker to detect changed metadata key");
        }
        catch (InvalidKeystoreException e)
        {
            // ok, expected
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }
    }

    @Test
    public void test5()
    {
        // create main keystore, backup main keystore, change main keystore -> check that backup keystore key is ok, re-register new main key
        // check that the new main keystore key has been re-registered

        // Firstly, re-install main keystore to register a well-known key
        encryptionKeysRegistry.unregisterKey(KeyProvider.ALIAS_METADATA);

        TestAlfrescoKeyStore keyStore = getKeyStore("main", "JCEKS", Collections.singletonMap(KeyProvider.ALIAS_METADATA, "metadata"),
                null, generateKeystoreName(), generateKeystoreName());
        keyStore.setKeysToValidate(new HashSet<>(Collections.singletonList("metadata")));
        createAndPopulateKeyStore(keyStore);

        try
        {
            keyStoreChecker.setMainKeyStore(keyStore);
            keyStoreChecker.validateKeyStores();
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception: " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }

        keyStore.backup();

        // check that the metadata key has been registered
        assertTrue("", encryptionKeysRegistry.isKeyRegistered("metadata"));

        // change the metadata key
        keyStore.changeKey(KeyProvider.ALIAS_METADATA, generateSecretKey("DESede"));

        try
        {
            // should detect changed metadata key and re-register it
            keyStoreChecker.validateKeyStores();
        }
        catch (InvalidKeystoreException e)
        {
            fail("Unexpected exception: " + e.getMessage());
        }
        catch (MissingKeyException e)
        {
            fail("Unexpected exception : " + e.getMessage());
        }

        // check that the new metadata key has been successfully re-registered by encrypting and decrypting some content with it
        assertTrue("", EncryptionKeysRegistry.KEY_STATUS.OK == encryptionKeysRegistry.checkKey(KeyProvider.ALIAS_METADATA,
                keyStore.getKey(KeyProvider.ALIAS_METADATA)));
    }

    private void createAndPopulateKeyStore(TestAlfrescoKeyStore keyStore)
    {
        KeyMap keyMap = new KeyMap();
        keyMap.setKey(KeyProvider.ALIAS_METADATA, generateSecretKey("DESede"));
        keyStore.create(keyMap, null);
    }

    private static class TestAlfrescoKeyStore extends AlfrescoKeyStoreImpl
    {
        public void create(KeyMap keys, KeyMap backupKeys)
        {
            this.keys = (keys != null ? keys : new KeyMap());
            this.backupKeys = (backupKeys != null ? backupKeys : new KeyMap());
            super.create();
        }

        void changeKey(String keyAlias, Key key)
        {
            keys.setKey(KeyProvider.ALIAS_METADATA, key);
        }
    }
}
